'''
Author: tomwoo tom.woo@outlook.com
Date: 2025-07-27 01:12:05
LastEditors: tomwoo tom.woo@outlook.com
LastEditTime: 2025-08-28 23:06:52
FilePath: /Weather_Forecast_Agent/main.py
Description: 这是默认设置,请设置`customMade`, 打开koroFileHeader查看配置 进行设置: https://github.com/OBKoro1/koro1FileHeader/wiki/%E9%85%8D%E7%BD%AE
'''

# 加载环境变量
from dotenv import load_dotenv, find_dotenv
_ = load_dotenv(find_dotenv())

import math
import os
from pathlib import Path
import shutil
from collections import OrderedDict
from datetime import datetime, timezone

import gradio as gr

from timing_tasks import set_interval
from forecast_tools.forecast import variables, get_current_datetime, try_to_generate_forecast, generate_animations
import forecast_tools.mcp_client as forecast_tools_mcp_client


timer_interval = 60 * 60 # unit: second
days_1 = 7
loc_1 = "广州"
lat_lon_1 = (+23.0, +113.0) # latitude, longitude
timezone_delta_1 = +8 # unit: hour
loc_2 = "地点"
lat_lon_2 = (0, 0) # (latitude, longitude)
timezone_delta_2 = 0 # unit: hour
llm_needed = True

js_file = "update-current-datetime.js"
dataset_path = "./dataset/"
dataset_file = "forecast_results.nc"
images_1_path = "./images_1/"
image_ext = ".gif"
images_1 = [None] * len(variables)
images_2 = [None] * len(variables)
config_path = "./config/"
loc_file = "locations.txt"
forecast_tools_mcp_server_path = "./forecast_tools/"
forecast_tools_mcp_server_file = "mcp_server.py"

# images_1 = ["images_1/msl.gif", "images_1/t2m.gif", "images_1/tcwv.gif"] # DEBUG


def cleanup_folder(folder_path: str) -> None:
    path = Path(folder_path)
    if path.exists():
        for item in path.iterdir():
            if item.is_dir():
                shutil.rmtree(item)  # Recursively delete subdirectories
            else:
                item.unlink()  # Delete files
        print(f"Folder '{folder_path}' has been cleaned up.")
    else:
        print(f"Folder '{folder_path}' does not exist.")

def init() -> None:
    global images_1

    print(f"UTC Date & Time: {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"Local Date & Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("init() running...")

    Path(dataset_path).mkdir(parents=True, exist_ok=True)
    Path(images_1_path).mkdir(parents=True, exist_ok=True)

    cleanup_folder(images_1_path)

    filenames = try_to_generate_forecast(os.path.join(dataset_path, dataset_file), 
                                         days_1, lat_lon_1, timezone_delta_1, 
                                         images_1_path, image_ext)

    for i, filename in enumerate(filenames):
        if filename:
            images_1[i] = filename

    status = client.connect(os.path.join(forecast_tools_mcp_server_path, forecast_tools_mcp_server_file))
    print(status)

def update() -> None:
    global images_1

    print(f"UTC Date & Time: {datetime.now(timezone.utc).strftime('%Y-%m-%d %H:%M:%S')}")
    print(f"Local Date & Time: {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}")
    print("update() running...")

    filenames = try_to_generate_forecast(os.path.join(dataset_path, dataset_file), 
                                         days_1, lat_lon_1, timezone_delta_1, 
                                         images_1_path, image_ext)

    for i, filename in enumerate(filenames):
        if filename:
            images_1[i] = filename

def update_current_datetime(timezone_delta: float) -> str:
    current_datetime = get_current_datetime(timezone_delta)
    return f"""## <center>世界标准时间：{current_datetime[0].strftime('%Y-%m-%d %H:%M:%S')}</center>
               ## <center>当地日期时间：{current_datetime[1].strftime('%Y-%m-%d %H:%M:%S')}</center>"""

def update_current_datetime_1() -> str:
    return update_current_datetime(timezone_delta_1)

def update_current_datetime_2() -> str:
    return update_current_datetime(timezone_delta_2)


client = forecast_tools_mcp_client.MCPClientWrapper()

init()

set_interval(update, timer_interval)

def markdown(days: int, loc: str, lat_lon: tuple[float, float], timezone_delta: float) -> str:
    if lat_lon[0] >= 0:
        lat = lat_lon[0]
        lat_orientation = 'N'
    else:
        lat = -lat_lon[0]
        lat_orientation = 'S'
    if lat_lon[1] >= 0:
        lon = lat_lon[1]
        lon_orientation = 'E'
    else:
        lon = -lat_lon[1]
        lon_orientation = 'W'

    return f"# <center>未来{days}天{loc}({lat}°{lat_orientation}, {lon}°{lon_orientation}, UTC{timezone_delta:+n}h)附近天气预报动画</center>"

def html(id: int, timezone_delta: float) -> str:
    return f"""
<div id="datetime-{id}a" style="text-align: center; font-size: 22px;"></div>
<div id="datetime-{id}b" style="text-align: center; font-size: 22px;" data-timezone-delta="{timezone_delta}"></div>
"""

js = Path(js_file).read_text()

loc_list = Path(os.path.join(config_path, loc_file)).read_text().split("\n")

default_values = OrderedDict({
    "dropdown": loc_list[0],
    "chatbot": [],
    "latitude_slider": 0,
    "longitude_slider": 0,
    "days_slider": 3,
    "interval_slider": 1000
})

with gr.Blocks(title="Weather Forecast Agent", head=f"<script>{js}</script>") as demo:
    def on_select(evt: gr.SelectData):
        global llm_needed
        match evt.value:
            case "地点":
                llm_needed = True
            case "经纬度坐标":
                llm_needed = False
            case _:
                pass

    def generating_animations_info(lat_lon, timezone_delta):
        gr.Info(f"纬度：{lat_lon[0]}°，经度：{lat_lon[1]}°，时差：{timezone_delta:+n}h", duration=5)
        gr.Info("正在生成天气预报动画...", duration=5)

    def generate_button_clicked_1(dropdown, latitude_slider, longitude_slider, days_slider, interval_slider):
        global loc_2, lat_lon_2, timezone_delta_2, images_2

        if llm_needed:
            loc_2 = dropdown

            system_message = "You are a helpful and harmless assistant. You should think step-by-step."
            user_message = f"请生成{dropdown}于{datetime.now(timezone.utc).date()}起{days_slider}天内的天气预报动画。动画的帧间间隔为{interval_slider}ms。"

            try:
                for result in client.query(user_message, system_message):
                    for msg in result["history"]:
                        try:
                            if "tool_call_generate_animations" == msg["metadata"]["id"] and "pending" == msg["metadata"]["status"]:
                                tool_calls = result["tool_calls"]
                                lat_lon_2 = tool_calls[0]["args"]["lat_lon"]
                                timezone_delta_2 = tool_calls[0]["args"]["timezone_delta"]
                                generating_animations_info(lat_lon_2, timezone_delta_2)
                                break
                        except Exception:
                            continue

                    try:
                        images_2 = result["tool_calls"][0]["result"]
                    except Exception:
                        pass

                    yield result["history"]
            except Exception as e:
                raise gr.Error(str(e), duration=5)
        else:
            loc_2 = "地点"
            lat_lon_2 = (latitude_slider, longitude_slider)
            timezone_delta_2 = math.copysign((abs(longitude_slider) + 7.5) // 15, longitude_slider)
            generating_animations_info(lat_lon_2, timezone_delta_2)

            try:
                images_2 = generate_animations(os.path.join(dataset_path, dataset_file), days_slider, 
                                               lat_lon_2, timezone_delta_2, 
                                               interval_slider)
            except Exception as e:
                raise gr.Error(str(e), duration=5)

    def generate_button_clicked_2(days_slider):
        global loc_2, lat_lon_2, timezone_delta_2, images_2
        return [gr.update(visible=True), markdown(days_slider, loc_2, lat_lon_2, timezone_delta_2), 
                html(2, timezone_delta_2), *images_2]

    with gr.Blocks() as part_1:
        gr.Markdown(markdown(days_1, loc_1, lat_lon_1, timezone_delta_1))
        # gr.Markdown(update_current_datetime_1, every=1)
        gr.HTML(html(1, timezone_delta_1))
        with gr.Row(equal_height=True):
            for i, var in enumerate(variables):
                gr.Image(images_1[i], label=var["ch_name"], show_label=True, interactive=False)

    with gr.Blocks() as part_2:
        gr.Markdown("# 你可以通过选择地点或经纬度坐标生成天气预报动画")
        with gr.Tab("地点") as tab_loc:
            dropdown = gr.Dropdown(loc_list, allow_custom_value=True, label="地点", show_label=True, info="你可以手动输入自选地点", interactive=True)
            chatbot = gr.Chatbot(type="messages", label=client.model, show_label=True, height=200, avatar_images=('👤', '🤖'), 
                                 placeholder="通过大模型确定用户指定地点的经纬度坐标及其今天的当地时间与世界标准时间的时差，并调用工具生成天气预报动画。")
        with gr.Tab("经纬度坐标") as tab_lat_lon:
            latitude_slider = gr.Slider(-90, +90, step=1, value=default_values["latitude_slider"], label="纬度(°)", show_label=True, info="北纬为正，南纬为负", interactive=True)
            longitude_slider = gr.Slider(-180, +180, step=1, value=default_values["longitude_slider"], label="经度(°)", show_label=True, info="东经为正，西经为负", interactive=True)
        with gr.Row(equal_height=True):
            days_slider = gr.Slider(1, 7, step=1, value=default_values["days_slider"], label="预报天数", show_label=True, interactive=True)
            interval_slider = gr.Slider(500, 5000, step=500, value=default_values["interval_slider"], label="帧间间隔(ms)", show_label=True, interactive=True)
            with gr.Column():
                restore_button = gr.Button("恢复默认值", variant="secondary")
                generate_button = gr.Button("生成天气预报动画", variant="primary")

    with gr.Blocks() as part_3:
        with gr.Column(visible=False) as part_3_col:
            markdown_2 = gr.Markdown()
            # gr.Markdown(update_current_datetime_2, every=1)
            html_2 = gr.HTML()
            with gr.Row(equal_height=True):
                gr_images_2 = []
                for i, var in enumerate(variables):
                    gr_images_2.append(gr.Image(format="gif", label=var["ch_name"], show_label=True, interactive=False))

    tab_loc.select(on_select, None, None)

    tab_lat_lon.select(on_select, None, None)

    restore_button.click(fn=lambda: list(default_values.values()),
                         outputs=[eval(key) for key in default_values.keys()])

    generate_button.click(generate_button_clicked_1, inputs=[dropdown, latitude_slider, longitude_slider, days_slider, interval_slider], outputs=[chatbot])\
                   .success(generate_button_clicked_2, inputs=[days_slider], outputs=[part_3_col, markdown_2, html_2, *gr_images_2])


if __name__ == "__main__":
    demo.queue().launch()

# end of file
